import json
import numpy as np


def calculate_mep(matched_position, true_position):
    """
    计算匹配误差像素 (MEP)。

    :param matched_position: 匹配位置 (x, y)。
    :param true_position: 真实位置 (x, y)。
    :return: 匹配误差像素 (MEP)。
    """
    mx, my = matched_position
    tx, ty = true_position
    return np.sqrt((mx - tx) ** 2 + (my - ty) ** 2)


def calculate_meps(matching_results):
    """
    计算所有 MEP。

    :param matching_results: 匹配结果列表。
    :return: 所有图像对的 MEP。
    """
    return [calculate_mep(result['matched_position'], result['true_position']) for result in matching_results]


def calculate_mae(meps, error_threshold):
    """
    计算平均绝对误差 (MAE)。

    :param meps: 各图像对的MEP值列表。
    :param error_threshold: 错误阈值。如果列表为空，返回 NaN。
    :return: 平均绝对误差 (MAE)。
    """
    filtered_meps = [mep for mep in meps if mep < error_threshold]
    if len(filtered_meps) == 0:
        return float('nan')
    return np.mean(filtered_meps)


def calculate_rmse(meps, error_threshold=None):
    """
    计算均方根误差 (RMSE)。可以选择带有错误阈值的计算。

    :param meps: 各图像对的MEP值列表。
    :param error_threshold: 错误阈值。如果列表为空，返回 NaN。
    :return: 均方根误差 (RMSE)。
    """
    filtered_meps = [mep for mep in meps if mep < error_threshold]
    if len(filtered_meps) == 0:
        return float('nan')

    return np.sqrt(np.mean(np.square(filtered_meps)))


def calculate_cmr(meps, error_threshold):
    """
    计算正确匹配率 (CMR)。

    :param meps: 各图像对的MEP值列表。
    :param error_threshold: 错误阈值。如果列表为空，返回 NaN。
    :return: 正确匹配率 (CMR)。
    """
    if len(meps) == 0:
        return float('nan')
    corr = sum(1 for mep in meps if mep < error_threshold)
    return corr / len(meps)


def evaluate_matching_results(matching_results_filepath, evaluation_results_filepath):
    """
    评估匹配结果，并将评估结果保存到 JSON 文件。

    :param matching_results_filepath: 匹配结果 JSON 文件路径。
    :param evaluation_results_filepath: 评估结果 JSON 文件路径。
    """
    with open(matching_results_filepath, 'r') as f:
        matching_results = json.load(f)

    meps = calculate_meps(matching_results)

    mae_3 = calculate_mae(meps, 3)
    mae_5 = calculate_mae(meps, 5)
    mae_10 = calculate_mae(meps, 10)

    rmse_3 = calculate_rmse(meps, 3)
    rmse_5 = calculate_rmse(meps, 5)
    rmse_10 = calculate_rmse(meps, 10)

    cmr_3 = calculate_cmr(meps, 3)
    cmr_5 = calculate_cmr(meps, 5)
    cmr_10 = calculate_cmr(meps, 10)

    evaluation_results = {
        "MAE": {
            "mae_3": mae_3,
            "mae_5": mae_5,
            "mae_10": mae_10
        },
        "RMSE": {
            "rmse_3": rmse_3,
            "rmse_5": rmse_5,
            "rmse_10": rmse_10
        },
        "CMR": {
            "cmr_3": cmr_3,
            "cmr_5": cmr_5,
            "cmr_10": cmr_10
        }
    }

    with open(evaluation_results_filepath, 'w') as f:
        json.dump(evaluation_results, f, indent=4)

    print(f"评估结果已保存到 {evaluation_results_filepath}")


if __name__ == "__main__":
    evaluate_matching_results('../data/optical_to_SAR/matching_results.json',
                              '../data/optical_to_SAR/evaluation_results.json')
